import torch
from torchvision import models, transforms
from torchvision.models import ResNet101_Weights
from PIL import Image

# 加载预训练模型
resnet = models.resnet101(weights=ResNet101_Weights.IMAGENET1K_V1)

# 定义预处理步骤
preprocess = transforms.Compose([
    transforms.Resize(256),  # 调整短边为256像素
    transforms.CenterCrop(224),  # 从中心裁剪224x224图像
    transforms.ToTensor(),  # 转换为张量
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])  # 归一化
])

# 打开和预处理图像
image = Image.open("../data/testimg/bobby.jpg")
image_t = preprocess(image)

# 批处理维度
batch_t = torch.unsqueeze(image_t, dim=0)

# 评估模式
resnet.eval()

# 前向传播
out = resnet(batch_t)

# 读取标签
with open('./imagenet_classes.txt') as f:
    labels = [line.strip() for line in f.readlines()]

# 获取预测结果
_, index = torch.max(out, 1)
percentage = torch.nn.functional.softmax(out, dim=1)[0] * 100
print(labels[index[0]], percentage[index[0]].item())

# 获取前5个预测结果
_, indices = torch.sort(out, descending=True)
result = [(labels[idx], percentage[idx].item()) for idx in indices[0][:5]]
print(result)

